{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Best Practices for Using Runners"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing installation: todd-ai 0.4.0\n",
      "Uninstalling todd-ai-0.4.0:\n",
      "  Successfully uninstalled todd-ai-0.4.0\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip uninstall -y todd_ai\n",
    "%pip install --no-build-isolation --extra-index-url https://pypi.org/simple .. > /dev/null\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:17:47,682 62058:140704275689088][patches.py:9 todd <module>] INFO: `ipdb` is installed. Using it for debugging.\n",
      "/Users/bytedance/.local/share/virtualenvs/todd-ARrcnwyq/lib/python3.11/site-packages/mmcv/__init__.py:20: UserWarning: On January 1, 2023, MMCV will release v2.0.0, in which it will remove components related to the training process and add a data transformation module. In addition, it will rename the package names mmcv to mmcv-lite and mmcv-full to mmcv. See https://github.com/open-mmlab/mmcv/blob/master/docs/en/compatibility.md for more details.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import pathlib\n",
    "import tempfile\n",
    "import time\n",
    "from pprint import pprint\n",
    "from typing import Any, NoReturn, TypedDict\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.data\n",
    "from torch import nn\n",
    "\n",
    "import todd\n",
    "from todd.runners import Memo"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@todd.ModelRegistry.register_()\n",
    "class RunnerModel(nn.Module):\n",
    "\n",
    "    def __init__(self) -> None:\n",
    "        super().__init__()\n",
    "        self._weight = torch.nn.Parameter(torch.tensor(0.0))\n",
    "\n",
    "    @property\n",
    "    def weight(self) -> torch.nn.Parameter:\n",
    "        return self._weight\n",
    "\n",
    "    def _forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        return x * self._weight\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        runner: todd.runners.BaseRunner,\n",
    "        batch,\n",
    "        memo: Memo,\n",
    "        *args,\n",
    "        **kwargs,\n",
    "    ) -> Memo:\n",
    "        log: dict[str, Any] | None = memo.get(\"log\")\n",
    "        y = self._forward(batch[\"x\"])\n",
    "        loss = F.l1_loss(y, batch[\"y\"])\n",
    "        memo[\"loss\"] = loss\n",
    "        if log is not None:\n",
    "            log[\"batch\"] = str(batch)\n",
    "            log[\"weight\"] = f\"{self._weight.item():.3f}\"\n",
    "            log[\"loss\"] = f\"{loss:.3f}\"\n",
    "        return memo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sample(TypedDict):\n",
    "    x: int\n",
    "    y: int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@todd.DatasetRegistry.register_()\n",
    "class RunnerDataset(torch.utils.data.Dataset[int]):\n",
    "\n",
    "    def __init__(self, n: int) -> None:\n",
    "        self._data = list(range(1, n + 1))\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self._data)\n",
    "\n",
    "    def __getitem__(self, index: int) -> Sample:\n",
    "        x = self._data[index]\n",
    "        return Sample(x=x, y=x * 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Batch(TypedDict):\n",
    "    x: torch.Tensor\n",
    "    y: torch.Tensor"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:17:49,375 62058:140704275689088][base.py:56 todd.Validator.validator __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1;36m/var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpjms97dk6\u001b[0m\n",
      "└── \u001b[1;36mvalidator\u001b[0m\n",
      "\n",
      "2 directories, 0 files\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type='Validator',\n",
    "    name='validator',\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.run()\n",
    "\n",
    "    !echo\n",
    "    !tree $work_dirs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:17:49,723 62058:140704275689088][base.py:56 todd.Validator.validator __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:17:49,728 62058:140704275689088][log.py:93 todd.Validator.validator after_run_iter] INFO: Iter [5/20] batch={'x': tensor([5]), 'y': tensor([10])} weight=0.000 loss=10.000\n",
      "[2024-03-14 12:17:49,731 62058:140704275689088][log.py:93 todd.Validator.validator after_run_iter] INFO: Iter [10/20] batch={'x': tensor([10]), 'y': tensor([20])} weight=0.000 loss=20.000\n",
      "[2024-03-14 12:17:49,734 62058:140704275689088][log.py:93 todd.Validator.validator after_run_iter] INFO: Iter [15/20] batch={'x': tensor([15]), 'y': tensor([30])} weight=0.000 loss=30.000\n",
      "[2024-03-14 12:17:49,736 62058:140704275689088][log.py:93 todd.Validator.validator after_run_iter] INFO: Iter [20/20] batch={'x': tensor([20]), 'y': tensor([40])} weight=0.000 loss=40.000\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1;36m/var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpv4q6cgrv\u001b[0m\n",
      "└── \u001b[1;36mvalidator\u001b[0m\n",
      "\n",
      "2 directories, 0 files\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type='Validator',\n",
    "    name='validator',\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    "    callbacks=[dict(type='LogCallback', interval=5)],\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.run()\n",
    "\n",
    "    !echo\n",
    "    !tree $work_dirs\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trainers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iteration Based"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:17:50,032 62058:140704275689088][base.py:56 todd.IterBasedTrainer.iter_based_trainer __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:17:50,035 62058:140704275689088][log.py:93 todd.IterBasedTrainer.iter_based_trainer after_run_iter] INFO: Iter [1/8] batch={'x': tensor([ 7, 10]), 'y': tensor([14, 20])} weight=0.000 loss=17.000\n",
      "[2024-03-14 12:17:50,037 62058:140704275689088][log.py:93 todd.IterBasedTrainer.iter_based_trainer after_run_iter] INFO: Iter [2/8] batch={'x': tensor([4, 6]), 'y': tensor([ 8, 12])} weight=0.000 loss=10.000\n",
      "[2024-03-14 12:17:50,038 62058:140704275689088][log.py:93 todd.IterBasedTrainer.iter_based_trainer after_run_iter] INFO: Iter [3/8] batch={'x': tensor([3, 8]), 'y': tensor([ 6, 16])} weight=0.000 loss=11.000\n",
      "[2024-03-14 12:17:50,039 62058:140704275689088][log.py:93 todd.IterBasedTrainer.iter_based_trainer after_run_iter] INFO: Iter [4/8] batch={'x': tensor([9, 2]), 'y': tensor([18,  4])} weight=0.000 loss=11.000\n",
      "[2024-03-14 12:17:50,041 62058:140704275689088][log.py:93 todd.IterBasedTrainer.iter_based_trainer after_run_iter] INFO: Iter [5/8] batch={'x': tensor([5, 1]), 'y': tensor([10,  2])} weight=0.000 loss=6.000\n",
      "[2024-03-14 12:17:50,043 62058:140704275689088][log.py:93 todd.IterBasedTrainer.iter_based_trainer after_run_iter] INFO: Iter [6/8] batch={'x': tensor([8, 1]), 'y': tensor([16,  2])} weight=0.000 loss=9.000\n",
      "[2024-03-14 12:17:50,044 62058:140704275689088][log.py:93 todd.IterBasedTrainer.iter_based_trainer after_run_iter] INFO: Iter [7/8] batch={'x': tensor([ 5, 10]), 'y': tensor([10, 20])} weight=0.000 loss=15.000\n",
      "[2024-03-14 12:17:50,046 62058:140704275689088][log.py:93 todd.IterBasedTrainer.iter_based_trainer after_run_iter] INFO: Iter [8/8] batch={'x': tensor([2, 4]), 'y': tensor([4, 8])} weight=0.000 loss=6.000\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"IterBasedTrainer\",\n",
    "    name=\"iter_based_trainer\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[dict(type=\"LogCallback\", interval=1)],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    iters=8,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Epoch Based"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:17:50,057 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.epoch_based_trainer __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:17:50,058 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.epoch_based_trainer before_run_epoch] INFO: Epoch [1/3]\n",
      "[2024-03-14 12:17:50,061 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [1/15] batch={'x': tensor([4, 1]), 'y': tensor([8, 2])} weight=0.000 loss=5.000\n",
      "[2024-03-14 12:17:50,064 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [2/15] batch={'x': tensor([8, 6]), 'y': tensor([16, 12])} weight=0.000 loss=14.000\n",
      "[2024-03-14 12:17:50,066 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [3/15] batch={'x': tensor([5, 9]), 'y': tensor([10, 18])} weight=0.000 loss=14.000\n",
      "[2024-03-14 12:17:50,068 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [4/15] batch={'x': tensor([ 7, 10]), 'y': tensor([14, 20])} weight=0.000 loss=17.000\n",
      "[2024-03-14 12:17:50,071 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [5/15] batch={'x': tensor([2, 3]), 'y': tensor([4, 6])} weight=0.000 loss=5.000\n",
      "[2024-03-14 12:17:50,073 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.epoch_based_trainer before_run_epoch] INFO: Epoch [2/3]\n",
      "[2024-03-14 12:17:50,075 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [6/15] batch={'x': tensor([7, 8]), 'y': tensor([14, 16])} weight=0.000 loss=15.000\n",
      "[2024-03-14 12:17:50,077 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [7/15] batch={'x': tensor([1, 6]), 'y': tensor([ 2, 12])} weight=0.000 loss=7.000\n",
      "[2024-03-14 12:17:50,079 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [8/15] batch={'x': tensor([2, 4]), 'y': tensor([4, 8])} weight=0.000 loss=6.000\n",
      "[2024-03-14 12:17:50,081 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [9/15] batch={'x': tensor([10,  9]), 'y': tensor([20, 18])} weight=0.000 loss=19.000\n",
      "[2024-03-14 12:17:50,082 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [10/15] batch={'x': tensor([3, 5]), 'y': tensor([ 6, 10])} weight=0.000 loss=8.000\n",
      "[2024-03-14 12:17:50,084 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.epoch_based_trainer before_run_epoch] INFO: Epoch [3/3]\n",
      "[2024-03-14 12:17:50,086 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [11/15] batch={'x': tensor([3, 6]), 'y': tensor([ 6, 12])} weight=0.000 loss=9.000\n",
      "[2024-03-14 12:17:50,088 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [12/15] batch={'x': tensor([2, 4]), 'y': tensor([4, 8])} weight=0.000 loss=6.000\n",
      "[2024-03-14 12:17:50,090 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [13/15] batch={'x': tensor([5, 8]), 'y': tensor([10, 16])} weight=0.000 loss=13.000\n",
      "[2024-03-14 12:17:50,091 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [14/15] batch={'x': tensor([10,  7]), 'y': tensor([20, 14])} weight=0.000 loss=17.000\n",
      "[2024-03-14 12:17:50,093 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.epoch_based_trainer after_run_iter] INFO: Iter [15/15] batch={'x': tensor([9, 1]), 'y': tensor([18,  2])} weight=0.000 loss=10.000\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"EpochBasedTrainer\",\n",
    "    name=\"epoch_based_trainer\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[dict(type=\"LogCallback\", interval=1)],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    epochs=3,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:17:50,216 62058:140704275689088][log.py:55 todd.Validator.log_callback init] INFO: \n",
      "platform: macOS-14.0\n",
      "nvidia_smi: None\n",
      "python_version: 3.11.7 (main, Dec  4 2023, 18:10:11) [Clang 15.0.0 (clang-1500.1.0.2.5)]\n",
      "pytorch_version: 2.0.1\n",
      "torchvision_version: 0.15.2\n",
      "opencv_version: 4.7.0\n",
      "todd_version: 0.4.0\n",
      "cuda_home: None\n",
      "git_commit_id: 0a7955a\n",
      "git_status: \n",
      "M todd/runners/callbacks/checkpoint.py\n",
      " M todd/runners/callbacks/composed.py\n",
      " M todd/runners/callbacks/git.py\n",
      " M todd/runners/callbacks/interval.py\n",
      " M todd/runners/callbacks/log.py\n",
      " M todd/runners/callbacks/lr.py\n",
      " M todd/runners/callbacks/monitor.py\n",
      " M todd/runners/callbacks/tensorboard.py\n",
      " M todd/runners/epoch_based_trainer.py\n",
      " M todd/runners/iter_based_trainer.py\n",
      " M todd/runners/strategies/base.py\n",
      " M todd/runners/strategies/ddp.py\n",
      " M todd/runners/strategies/fsdp.py\n",
      " M todd/runners/utils.py\n",
      " M todd/utils/__init__.py\n",
      " M todd/utils/mixins.py\n",
      "?? todd/utils/constants.py\n",
      "\u001b[2m[2024-03-14 12:17:50,217 62058:140704275689088][base.py:56 todd.Validator.log_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:17:50,221 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [5/20] batch={'x': tensor([5]), 'y': tensor([10])} weight=0.000 loss=10.000\n",
      "[2024-03-14 12:17:50,240 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [10/20] batch={'x': tensor([10]), 'y': tensor([20])} weight=0.000 loss=20.000\n",
      "[2024-03-14 12:17:50,247 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [15/20] batch={'x': tensor([15]), 'y': tensor([30])} weight=0.000 loss=30.000\n",
      "[2024-03-14 12:17:50,262 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [20/20] batch={'x': tensor([20]), 'y': tensor([40])} weight=0.000 loss=40.000\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"Validator\",\n",
    "    name=\"log_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    "    callbacks=[\n",
    "        dict(\n",
    "            type=\"LogCallback\",\n",
    "            interval=5,\n",
    "            collect_env=dict(verbose=False),\n",
    "        ),\n",
    "    ],\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:17:50,291 62058:140704275689088][base.py:56 todd.Validator.log_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:17:50,294 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [5/20] batch={'x': tensor([5]), 'y': tensor([10])} weight=0.000 loss=10.000\n",
      "[2024-03-14 12:17:50,296 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [10/20] batch={'x': tensor([10]), 'y': tensor([20])} weight=0.000 loss=20.000\n",
      "[2024-03-14 12:17:50,299 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [15/20] batch={'x': tensor([15]), 'y': tensor([30])} weight=0.000 loss=30.000\n",
      "[2024-03-14 12:17:50,301 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [20/20] batch={'x': tensor([20]), 'y': tensor([40])} weight=0.000 loss=40.000\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1;36m/var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpypx13e0i\u001b[0m\n",
      "└── \u001b[1;36mlog_callback\u001b[0m\n",
      "    └── 2024-03-14T12-17-50_290843-08-00.log\n",
      "\n",
      "2 directories, 1 file\n",
      "\n",
      "[2024-03-14 12:17:50,291 62058:140704275689088][base.py:56 todd.Validator.log_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\n",
      "[2024-03-14 12:17:50,294 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [5/20] batch={'x': tensor([5]), 'y': tensor([10])} weight=0.000 loss=10.000\n",
      "[2024-03-14 12:17:50,296 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [10/20] batch={'x': tensor([10]), 'y': tensor([20])} weight=0.000 loss=20.000\n",
      "[2024-03-14 12:17:50,299 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [15/20] batch={'x': tensor([15]), 'y': tensor([30])} weight=0.000 loss=30.000\n",
      "[2024-03-14 12:17:50,301 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [20/20] batch={'x': tensor([20]), 'y': tensor([40])} weight=0.000 loss=40.000\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type='Validator',\n",
    "    name='log_callback',\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    "    callbacks=[\n",
    "        dict(\n",
    "            type='LogCallback',\n",
    "            interval=5,\n",
    "            with_file_handler=True,\n",
    "        ),\n",
    "    ],\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.run()\n",
    "\n",
    "    !echo\n",
    "    !tree {work_dirs}\n",
    "    !echo\n",
    "    !cat {work_dirs}/log_callback/*.log\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:17:50,870 62058:140704275689088][base.py:56 todd.Validator.log_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:17:51,391 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [5/20] ETA 0:00:01 batch={'x': tensor([5]), 'y': tensor([10])} weight=0.000 loss=10.000\n",
      "[2024-03-14 12:17:51,912 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [10/20] ETA 0:00:01 batch={'x': tensor([10]), 'y': tensor([20])} weight=0.000 loss=20.000\n",
      "[2024-03-14 12:17:52,431 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [15/20] ETA 0:00:00 batch={'x': tensor([15]), 'y': tensor([30])} weight=0.000 loss=30.000\n",
      "[2024-03-14 12:17:52,945 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [20/20] ETA 0:00:00 batch={'x': tensor([20]), 'y': tensor([40])} weight=0.000 loss=40.000\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"Validator\",\n",
    "    name=\"log_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    "    callbacks=[\n",
    "        dict(\n",
    "            type=\"LogCallback\",\n",
    "            interval=5,\n",
    "            eta=dict(type=\"AverageETA\"),\n",
    "        ),\n",
    "    ],\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.strategy.module.register_forward_hook(\n",
    "        lambda *args, **kwargs: time.sleep(0.1)\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:17:52,955 62058:140704275689088][base.py:56 todd.Validator.log_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:17:54,468 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [5/20] ETA 0:00:04 batch={'x': tensor([5]), 'y': tensor([10])} weight=0.000 loss=10.000\n",
      "[2024-03-14 12:17:58,481 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [10/20] ETA 0:00:05 batch={'x': tensor([10]), 'y': tensor([20])} weight=0.000 loss=20.000\n",
      "[2024-03-14 12:18:03,499 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [15/20] ETA 0:00:03 batch={'x': tensor([15]), 'y': tensor([30])} weight=0.000 loss=30.000\n",
      "[2024-03-14 12:18:08,518 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [20/20] ETA 0:00:00 batch={'x': tensor([20]), 'y': tensor([40])} weight=0.000 loss=40.000\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"Validator\",\n",
    "    name=\"log_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    "    callbacks=[\n",
    "        dict(\n",
    "            type=\"LogCallback\",\n",
    "            interval=5,\n",
    "            eta=dict(type=\"EMA_ETA\", ema=dict(decay=0.2)),\n",
    "        ),\n",
    "    ],\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.strategy.module.register_forward_hook(\n",
    "        lambda *args, **kwargs: time.sleep(0.1 * min(10, runner.iter_))\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:18:08,647 62058:140704275689088][log.py:55 todd.Validator.log_callback init] INFO: \n",
      "platform: macOS-14.0\n",
      "nvidia_smi: None\n",
      "python_version: 3.11.7 (main, Dec  4 2023, 18:10:11) [Clang 15.0.0 (clang-1500.1.0.2.5)]\n",
      "pytorch_version: 2.0.1\n",
      "torchvision_version: 0.15.2\n",
      "opencv_version: 4.7.0\n",
      "todd_version: 0.4.0\n",
      "cuda_home: None\n",
      "git_commit_id: 0a7955a\n",
      "git_status: \n",
      "M todd/runners/callbacks/checkpoint.py\n",
      " M todd/runners/callbacks/composed.py\n",
      " M todd/runners/callbacks/git.py\n",
      " M todd/runners/callbacks/interval.py\n",
      " M todd/runners/callbacks/log.py\n",
      " M todd/runners/callbacks/lr.py\n",
      " M todd/runners/callbacks/monitor.py\n",
      " M todd/runners/callbacks/tensorboard.py\n",
      " M todd/runners/epoch_based_trainer.py\n",
      " M todd/runners/iter_based_trainer.py\n",
      " M todd/runners/strategies/base.py\n",
      " M todd/runners/strategies/ddp.py\n",
      " M todd/runners/strategies/fsdp.py\n",
      " M todd/runners/utils.py\n",
      " M todd/utils/__init__.py\n",
      " M todd/utils/mixins.py\n",
      "?? todd/utils/constants.py\n",
      "\u001b[2m[2024-03-14 12:18:08,648 62058:140704275689088][base.py:56 todd.Validator.log_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:08,653 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [5/20] ETA 0:00:00 batch={'x': tensor([5]), 'y': tensor([10])} weight=0.000 loss=10.000\n",
      "[2024-03-14 12:18:08,655 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [10/20] ETA 0:00:00 batch={'x': tensor([10]), 'y': tensor([20])} weight=0.000 loss=20.000\n",
      "[2024-03-14 12:18:08,658 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [15/20] ETA 0:00:00 batch={'x': tensor([15]), 'y': tensor([30])} weight=0.000 loss=30.000\n",
      "[2024-03-14 12:18:08,660 62058:140704275689088][log.py:93 todd.Validator.log_callback after_run_iter] INFO: Iter [20/20] ETA 0:00:00 batch={'x': tensor([20]), 'y': tensor([40])} weight=0.000 loss=40.000\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"Validator\",\n",
    "    name=\"log_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    "    callbacks=[\n",
    "        dict(\n",
    "            type=\"LogCallback\",\n",
    "            interval=5,\n",
    "            collect_env=dict(verbose=False),\n",
    "            with_file_handler=True,\n",
    "            eta=dict(type=\"AverageETA\"),\n",
    "        ),\n",
    "    ],\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:18:08,732 62058:140704275689088][git.py:41 todd.Validator.git_callback init] INFO: Saving git diff to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpotx2rbek/git_callback/git_diff_2024-03-14T12-18-08_732484-08-00.log\n",
      "\u001b[2m[2024-03-14 12:18:08,735 62058:140704275689088][base.py:56 todd.Validator.git_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "diff --git a/todd/runners/callbacks/checkpoint.py b/todd/runners/callbacks/checkpoint.py\n",
      "index 09548a2..caba779 100644\n",
      "--- a/todd/runners/callbacks/checkpoint.py\n",
      "+++ b/todd/runners/callbacks/checkpoint.py\n",
      "@@ -36,27 +36,27 @@ class CheckpointCallback(IntervalMixin, BaseCallback):\n",
      " \n",
      "     def init(self, *args, **kwargs) -> None:\n",
      "         super().init(*args, **kwargs)\n",
      "-        self._checkpoint_dir = self._runner.work_dir / 'checkpoints'\n",
      "+        self._checkpoint_dir = self.runner.work_dir / 'checkpoints'\n",
      "         self._latest_checkpoint_dir = self._checkpoint_dir / 'latest'\n",
      " \n",
      "         self._checkpoint_dir.mkdir(parents=True, exist_ok=True)\n",
      " \n",
      "-        if self._runner._auto_resume and self._latest_checkpoint_dir.exists():\n",
      "+        if self.runner._auto_resume and self._latest_checkpoint_dir.exists():\n",
      "             load_from = self._latest_checkpoint_dir\n",
      "-        elif self._runner.load_from is not None:\n",
      "-            load_from = pathlib.Path(self._runner.load_from)\n",
      "+        elif self.runner.load_from is not None:\n",
      "+            load_from = pathlib.Path(self.runner.load_from)\n",
      "             assert load_from.exists()\n",
      "         else:\n",
      "             load_from = None\n",
      " \n",
      "         if load_from is not None:\n",
      "             if get_rank() == 0:\n",
      "-                self._runner.logger.info(\"Loading from %s\", load_from)\n",
      "+                self.runner.logger.info(\"Loading from %s\", load_from)\n",
      "             state_dict = {\n",
      "                 f.stem: torch.load(f, 'cpu')\n",
      "                 for f in load_from.glob('*.pth')\n",
      "             }\n",
      "-            self._runner.load_state_dict(state_dict, **self._load_state_dict)\n",
      "+            self.runner.load_state_dict(state_dict, **self._load_state_dict)\n",
      " \n",
      "     @property\n",
      "     def checkpoint_dir(self) -> pathlib.Path:\n",
      "@@ -71,13 +71,13 @@ class CheckpointCallback(IntervalMixin, BaseCallback):\n",
      " \n",
      "     def _save(self, name: str) -> None:\n",
      "         # for FSDP, all ranks should call state dict\n",
      "-        state_dict = self._runner.state_dict(**self._state_dict)\n",
      "+        state_dict = self.runner.state_dict(**self._state_dict)\n",
      " \n",
      "         if get_rank() != 0:\n",
      "             return\n",
      "         work_dir = self._work_dir(name)\n",
      "         work_dir.mkdir(parents=True, exist_ok=True)\n",
      "-        self._runner.logger.info(\"Saving state dict to %s\", work_dir)\n",
      "+        self.runner.logger.info(\"Saving state dict to %s\", work_dir)\n",
      "         for k, v in state_dict.items():\n",
      "             torch.save(v, work_dir / f'{k}.pth')\n",
      " \n",
      "@@ -88,7 +88,7 @@ class CheckpointCallback(IntervalMixin, BaseCallback):\n",
      "     def after_run_iter(self, batch, memo: Memo) -> None:\n",
      "         super().after_run_iter(batch, memo)\n",
      "         if self._should_run_iter():\n",
      "-            self._save(f'iter_{self._runner.iter_}')\n",
      "+            self._save(f'iter_{self.runner.iter_}')\n",
      " \n",
      "     def after_run_epoch(self, epoch_memo: Memo, memo: Memo) -> None:\n",
      "         super().after_run_epoch(epoch_memo, memo)\n",
      "diff --git a/todd/runners/callbacks/composed.py b/todd/runners/callbacks/composed.py\n",
      "index a1c1638..83bec29 100644\n",
      "--- a/todd/runners/callbacks/composed.py\n",
      "+++ b/todd/runners/callbacks/composed.py\n",
      "@@ -21,7 +21,7 @@ class ComposedCallback(BaseCallback):\n",
      "         super().__init__(*args, **kwargs)\n",
      "         priorities = [c.pop('priority', dict()) for c in callbacks]\n",
      "         queue = [\n",
      "-            CallbackRegistry.build(c, runner=self._runner) for c in callbacks\n",
      "+            CallbackRegistry.build(c, runner=self.runner) for c in callbacks\n",
      "         ]\n",
      "         self._priority_queue: PriorityQueue[KT, BaseCallback] = \\\n",
      "             PriorityQueue(priorities, queue)\n",
      "diff --git a/todd/runners/callbacks/git.py b/todd/runners/callbacks/git.py\n",
      "index ce642ac..ab90b0f 100644\n",
      "--- a/todd/runners/callbacks/git.py\n",
      "+++ b/todd/runners/callbacks/git.py\n",
      "@@ -33,10 +33,10 @@ class GitCallback(BaseCallback):\n",
      "                 diff = subprocess_run(args_)\n",
      "             except subprocess.CalledProcessError as e:\n",
      "                 diff = str(e)\n",
      "-                self._runner.logger.error(e)\n",
      "+                self.runner.logger.error(e)\n",
      "             else:\n",
      "                 file = (\n",
      "-                    self._runner.work_dir / f'git_diff_{get_timestamp()}.log'\n",
      "+                    self.runner.work_dir / f'git_diff_{get_timestamp()}.log'\n",
      "                 )\n",
      "-                self._runner.logger.info('Saving git diff to %s', file)\n",
      "+                self.runner.logger.info('Saving git diff to %s', file)\n",
      "                 file.write_text(diff)\n",
      "diff --git a/todd/runners/callbacks/interval.py b/todd/runners/callbacks/interval.py\n",
      "index 8eb9625..489dd21 100644\n",
      "--- a/todd/runners/callbacks/interval.py\n",
      "+++ b/todd/runners/callbacks/interval.py\n",
      "@@ -22,7 +22,7 @@ class IntervalMixin(BaseCallback):\n",
      "         return self._interval > 0 and step % self._interval == 0\n",
      " \n",
      "     def _should_run_iter(self) -> bool:\n",
      "-        return not self._by_epoch and self.__should_run(self._runner.iter_)\n",
      "+        return not self._by_epoch and self.__should_run(self.runner.iter_)\n",
      " \n",
      "     def _should_run_epoch(self) -> bool:\n",
      "         return (\n",
      "diff --git a/todd/runners/callbacks/log.py b/todd/runners/callbacks/log.py\n",
      "index 8f884f5..263e6b4 100644\n",
      "--- a/todd/runners/callbacks/log.py\n",
      "+++ b/todd/runners/callbacks/log.py\n",
      "@@ -43,24 +43,24 @@ class LogCallback(IntervalMixin, BaseCallback):\n",
      "         if get_rank() > 0:\n",
      "             return\n",
      "         if self._with_file_handler:\n",
      "-            file = self._runner.work_dir / f'{get_timestamp()}.log'\n",
      "+            file = self.runner.work_dir / f'{get_timestamp()}.log'\n",
      "             handler = logging.FileHandler(file)\n",
      "             handler.setFormatter(Formatter())\n",
      "-            self._runner.logger.addHandler(handler)\n",
      "+            self.runner.logger.addHandler(handler)\n",
      "         if self._collect_env is not None:\n",
      "             from ...base import (  # noqa: E501 pylint: disable=import-outside-toplevel\n",
      "                 collect_env,\n",
      "             )\n",
      "             env = collect_env(**self._collect_env)\n",
      "-            self._runner.logger.info(env)\n",
      "+            self.runner.logger.info(env)\n",
      " \n",
      "     def before_run(self, memo: Memo) -> None:\n",
      "         super().before_run(memo)\n",
      "         self._eta: BaseETA | None = (\n",
      "             None if self._eta_config is None else ETARegistry.build(\n",
      "                 self._eta_config,\n",
      "-                start=self._runner.iter_ - 1,\n",
      "-                end=self._runner.iters,\n",
      "+                start=self.runner.iter_ - 1,\n",
      "+                end=self.runner.iters,\n",
      "             )\n",
      "         )\n",
      " \n",
      "@@ -73,10 +73,10 @@ class LogCallback(IntervalMixin, BaseCallback):\n",
      "         super().after_run_iter(batch, memo)\n",
      "         if 'log' not in memo:\n",
      "             return\n",
      "-        prefix = f\"Iter [{self._runner.iter_}/{self._runner.iters}] \"\n",
      "+        prefix = f\"Iter [{self.runner.iter_}/{self.runner.iters}] \"\n",
      " \n",
      "         if self._eta is not None:\n",
      "-            eta = self._eta(self._runner.iter_)\n",
      "+            eta = self._eta(self.runner.iter_)\n",
      "             eta = round(eta)\n",
      "             prefix += f\"ETA {str(datetime.timedelta(seconds=eta))} \"\n",
      " \n",
      "@@ -90,7 +90,7 @@ class LogCallback(IntervalMixin, BaseCallback):\n",
      " \n",
      "         log: dict[str, Any] = memo.pop('log')\n",
      "         message = ' '.join(f'{k}={v}' for k, v in log.items() if v is not None)\n",
      "-        self._runner.logger.info(prefix + message)\n",
      "+        self.runner.logger.info(prefix + message)\n",
      " \n",
      "     def before_run_epoch(self, epoch_memo: Memo, memo: Memo) -> None:\n",
      "         super().before_run_epoch(epoch_memo, memo)\n",
      "diff --git a/todd/runners/callbacks/lr.py b/todd/runners/callbacks/lr.py\n",
      "index 977558b..6a597c9 100644\n",
      "--- a/todd/runners/callbacks/lr.py\n",
      "+++ b/todd/runners/callbacks/lr.py\n",
      "@@ -26,7 +26,7 @@ class LRScheduleCallback(IntervalMixin, BaseCallback):\n",
      "         **kwargs,\n",
      "     ) -> None:\n",
      "         super().__init__(*args, interval=interval, **kwargs)\n",
      "-        assert isinstance(self._runner, Trainer)\n",
      "+        assert isinstance(self.runner, Trainer)\n",
      "         self._lr_scheduler_config = lr_scheduler\n",
      " \n",
      "     def init(self, *args, **kwargs) -> None:\n",
      "@@ -34,7 +34,7 @@ class LRScheduleCallback(IntervalMixin, BaseCallback):\n",
      "         self._build_lr_scheduler()\n",
      " \n",
      "     def _build_lr_scheduler(self) -> None:\n",
      "-        runner = cast(Trainer, self._runner)\n",
      "+        runner = cast(Trainer, self.runner)\n",
      "         self._lr_scheduler: torch.optim.lr_scheduler.LRScheduler = \\\n",
      "             LRSchedulerRegistry.build(\n",
      "                 self._lr_scheduler_config,\n",
      "@@ -75,11 +75,11 @@ class LRScaleCallback(BaseCallback):\n",
      " \n",
      "     def __init__(self, *args, lr_scaler: Config, **kwargs) -> None:\n",
      "         super().__init__(*args, **kwargs)\n",
      "-        assert isinstance(self._runner, Trainer)\n",
      "+        assert isinstance(self.runner, Trainer)\n",
      "         self._lr_scaler_config = lr_scaler\n",
      " \n",
      "     def _scale_lr(self, config: Config) -> None:\n",
      "-        runner = cast(Trainer, self._runner)\n",
      "+        runner = cast(Trainer, self.runner)\n",
      "         assert runner.dataloader.batch_size is not None\n",
      "         base_batch_size = config.base_batch_size\n",
      "         batch_size = get_world_size() * runner.dataloader.batch_size\n",
      "diff --git a/todd/runners/callbacks/monitor.py b/todd/runners/callbacks/monitor.py\n",
      "index 98aabf2..0895477 100644\n",
      "--- a/todd/runners/callbacks/monitor.py\n",
      "+++ b/todd/runners/callbacks/monitor.py\n",
      "@@ -37,8 +37,8 @@ class MonitorCallback(BaseCallback):\n",
      "     ) -> None:\n",
      "         super().run_iter_context(exit_stack, batch, memo)\n",
      "         context = Context(\n",
      "-            self._runner.logger,\n",
      "-            iter_=self._runner.iter_,\n",
      "+            self.runner.logger,\n",
      "+            iter_=self.runner.iter_,\n",
      "             batch=batch,\n",
      "             memo=memo,\n",
      "         )\n",
      "diff --git a/todd/runners/callbacks/tensorboard.py b/todd/runners/callbacks/tensorboard.py\n",
      "index 65eb9bc..804f907 100644\n",
      "--- a/todd/runners/callbacks/tensorboard.py\n",
      "+++ b/todd/runners/callbacks/tensorboard.py\n",
      "@@ -31,7 +31,7 @@ class TensorBoardCallback(IntervalMixin, BaseCallback):\n",
      "         super().init(*args, **kwargs)\n",
      "         if get_rank() > 0:\n",
      "             return\n",
      "-        log_dir = self._runner.work_dir / 'tensorboard'\n",
      "+        log_dir = self.runner.work_dir / 'tensorboard'\n",
      "         self._summary_writer = SummaryWriter(\n",
      "             log_dir,\n",
      "             **self._summary_writer_config,\n",
      "diff --git a/todd/runners/epoch_based_trainer.py b/todd/runners/epoch_based_trainer.py\n",
      "index bc46346..8358d00 100644\n",
      "--- a/todd/runners/epoch_based_trainer.py\n",
      "+++ b/todd/runners/epoch_based_trainer.py\n",
      "@@ -18,7 +18,7 @@ T = TypeVar('T', bound=nn.Module)\n",
      " \n",
      " \n",
      " @RunnerRegistry.register_()\n",
      "-class EpochBasedTrainer(Trainer):\n",
      "+class EpochBasedTrainer(Trainer[T]):\n",
      " \n",
      "     def __init__(self, *args, epochs: int, **kwargs) -> None:\n",
      "         super().__init__(*args, **kwargs)\n",
      "diff --git a/todd/runners/iter_based_trainer.py b/todd/runners/iter_based_trainer.py\n",
      "index 107c894..7a212fa 100644\n",
      "--- a/todd/runners/iter_based_trainer.py\n",
      "+++ b/todd/runners/iter_based_trainer.py\n",
      "@@ -16,7 +16,7 @@ T = TypeVar('T', bound=nn.Module)\n",
      " \n",
      " \n",
      " @RunnerRegistry.register_()\n",
      "-class IterBasedTrainer(Trainer):\n",
      "+class IterBasedTrainer(Trainer[T]):\n",
      " \n",
      "     def __init__(self, *args, iters: int, **kwargs) -> None:\n",
      "         super().__init__(*args, **kwargs)\n",
      "diff --git a/todd/runners/strategies/base.py b/todd/runners/strategies/base.py\n",
      "index a97d30f..d6c2a74 100644\n",
      "--- a/todd/runners/strategies/base.py\n",
      "+++ b/todd/runners/strategies/base.py\n",
      "@@ -2,7 +2,7 @@ __all__ = [\n",
      "     'BaseStrategy',\n",
      " ]\n",
      " \n",
      "-from typing import Any, Generic, Mapping, TypeVar, cast\n",
      "+from typing import Any, Mapping, TypeVar, cast\n",
      " \n",
      " import torch\n",
      " from torch import nn\n",
      "@@ -15,7 +15,7 @@ T = TypeVar('T', bound=nn.Module)\n",
      " \n",
      " \n",
      " @StrategyRegistry.register_()\n",
      "-class BaseStrategy(RunnerHolderMixin, StateDictMixin, Generic[T]):\n",
      "+class BaseStrategy(RunnerHolderMixin[T], StateDictMixin):\n",
      " \n",
      "     def __init__(\n",
      "         self,\n",
      "@@ -45,7 +45,7 @@ class BaseStrategy(RunnerHolderMixin, StateDictMixin, Generic[T]):\n",
      " \n",
      "     @property\n",
      "     def module(self) -> nn.Module:\n",
      "-        return self._runner.model\n",
      "+        return self.runner.model\n",
      " \n",
      "     def model_state_dict(self, *args, **kwargs) -> dict[str, Any]:\n",
      "         return self.module.state_dict(*args, **kwargs)\n",
      "@@ -62,7 +62,7 @@ class BaseStrategy(RunnerHolderMixin, StateDictMixin, Generic[T]):\n",
      "             **kwargs,\n",
      "         )\n",
      "         if get_rank() == 0:\n",
      "-            self._runner.logger.info(incompatible_keys)\n",
      "+            self.runner.logger.info(incompatible_keys)\n",
      " \n",
      "     def load_model_from(\n",
      "         self,\n",
      "@@ -77,7 +77,7 @@ class BaseStrategy(RunnerHolderMixin, StateDictMixin, Generic[T]):\n",
      "         model_state_dict = dict()\n",
      "         for f_ in f_list:\n",
      "             if get_rank() == 0:\n",
      "-                self._runner.logger.info(\"Loading model from %s\", f_)\n",
      "+                self.runner.logger.info(\"Loading model from %s\", f_)\n",
      "             model_state_dict.update(torch.load(f_, 'cpu'))\n",
      "         self.load_model_state_dict(model_state_dict, *args, **kwargs)\n",
      " \n",
      "diff --git a/todd/runners/strategies/ddp.py b/todd/runners/strategies/ddp.py\n",
      "index 3367c2a..e2e9fd6 100644\n",
      "--- a/todd/runners/strategies/ddp.py\n",
      "+++ b/todd/runners/strategies/ddp.py\n",
      "@@ -23,4 +23,4 @@ class DDPStrategy(CUDAStrategy[T]):\n",
      " \n",
      "     @property\n",
      "     def module(self) -> nn.Module:\n",
      "-        return self._runner.model.module\n",
      "+        return self.runner.model.module\n",
      "diff --git a/todd/runners/strategies/fsdp.py b/todd/runners/strategies/fsdp.py\n",
      "index fa71105..42eb809 100644\n",
      "--- a/todd/runners/strategies/fsdp.py\n",
      "+++ b/todd/runners/strategies/fsdp.py\n",
      "@@ -26,13 +26,13 @@ class FSDPStrategy(CUDAStrategy[T]):\n",
      " \n",
      "     @property\n",
      "     def module(self) -> nn.Module:\n",
      "-        return self._runner.model.module\n",
      "+        return self.runner.model.module\n",
      " \n",
      "     def build_optimizer(self, config: Config) -> torch.optim.Optimizer:\n",
      "-        return OptimizerRegistry.build(config, model=self._runner.model)\n",
      "+        return OptimizerRegistry.build(config, model=self.runner.model)\n",
      " \n",
      "     def model_state_dict(self, *args, **kwargs) -> dict[str, Any]:\n",
      "-        return self._runner.model.state_dict(*args, **kwargs)\n",
      "+        return self.runner.model.state_dict(*args, **kwargs)\n",
      " \n",
      "     def load_model_state_dict(\n",
      "         self,\n",
      "@@ -40,7 +40,7 @@ class FSDPStrategy(CUDAStrategy[T]):\n",
      "         *args,\n",
      "         **kwargs,\n",
      "     ) -> None:\n",
      "-        self._runner.model.load_state_dict(state_dict, *args, **kwargs)\n",
      "+        self.runner.model.load_state_dict(state_dict, *args, **kwargs)\n",
      " \n",
      "     def optim_state_dict(\n",
      "         self,\n",
      "diff --git a/todd/runners/utils.py b/todd/runners/utils.py\n",
      "index 0f25783..9c72605 100644\n",
      "--- a/todd/runners/utils.py\n",
      "+++ b/todd/runners/utils.py\n",
      "@@ -2,42 +2,45 @@ __all__ = [\n",
      "     'RunnerHolderMixin',\n",
      " ]\n",
      " \n",
      "-import weakref\n",
      "-from typing import cast\n",
      "+from typing import TypeVar\n",
      " \n",
      "+from torch import nn\n",
      "+\n",
      "+from ..utils import HolderMixin\n",
      " from .base import BaseRunner\n",
      " from .epoch_based_trainer import EpochBasedTrainer\n",
      " from .iter_based_trainer import IterBasedTrainer\n",
      " from .trainer import Trainer\n",
      " from .validator import Validator\n",
      " \n",
      "+T = TypeVar('T', bound=nn.Module)\n",
      "+\n",
      " \n",
      "-class RunnerHolderMixin:\n",
      "+class RunnerHolderMixin(HolderMixin[BaseRunner[T]]):\n",
      " \n",
      "-    def __init__(self, *args, runner: BaseRunner, **kwargs) -> None:\n",
      "-        super().__init__(*args, **kwargs)\n",
      "-        runner_proxy = (\n",
      "-            runner if isinstance(runner, weakref.ProxyTypes) else\n",
      "-            weakref.proxy(runner)\n",
      "-        )\n",
      "-        self._runner = cast(BaseRunner, runner_proxy)\n",
      "+    def __init__(self, *args, runner: BaseRunner[T], **kwargs) -> None:\n",
      "+        super().__init__(*args, instance=runner, **kwargs)\n",
      "+\n",
      "+    @property\n",
      "+    def runner(self) -> BaseRunner[T]:\n",
      "+        return self._instance\n",
      " \n",
      "     @property\n",
      "-    def trainer(self) -> Trainer:\n",
      "-        assert isinstance(self._runner, Trainer)\n",
      "-        return self._runner\n",
      "+    def trainer(self) -> Trainer[T]:\n",
      "+        assert isinstance(self._instance, Trainer)\n",
      "+        return self._instance\n",
      " \n",
      "     @property\n",
      "-    def validator(self) -> Validator:\n",
      "-        assert isinstance(self._runner, Validator)\n",
      "-        return self._runner\n",
      "+    def validator(self) -> Validator[T]:\n",
      "+        assert isinstance(self._instance, Validator)\n",
      "+        return self._instance\n",
      " \n",
      "     @property\n",
      "-    def iter_based_trainer(self) -> IterBasedTrainer:\n",
      "-        assert isinstance(self._runner, IterBasedTrainer)\n",
      "-        return self._runner\n",
      "+    def iter_based_trainer(self) -> IterBasedTrainer[T]:\n",
      "+        assert isinstance(self._instance, IterBasedTrainer)\n",
      "+        return self._instance\n",
      " \n",
      "     @property\n",
      "-    def epoch_based_trainer(self) -> EpochBasedTrainer:\n",
      "-        assert isinstance(self._runner, EpochBasedTrainer)\n",
      "-        return self._runner\n",
      "+    def epoch_based_trainer(self) -> EpochBasedTrainer[T]:\n",
      "+        assert isinstance(self._instance, EpochBasedTrainer)\n",
      "+        return self._instance\n",
      "diff --git a/todd/utils/__init__.py b/todd/utils/__init__.py\n",
      "index 85d1d65..199b7df 100644\n",
      "--- a/todd/utils/__init__.py\n",
      "+++ b/todd/utils/__init__.py\n",
      "@@ -1,3 +1,4 @@\n",
      "+from .constants import *\n",
      " from .enums import *\n",
      " from .generic_tensors import *\n",
      " from .metas import *\n",
      "diff --git a/todd/utils/mixins.py b/todd/utils/mixins.py\n",
      "index 9574ea1..3e9a3cc 100644\n",
      "--- a/todd/utils/mixins.py\n",
      "+++ b/todd/utils/mixins.py\n",
      "@@ -1,8 +1,12 @@\n",
      " __all__ = [\n",
      "     'StateDictMixin',\n",
      "+    'HolderMixin',\n",
      " ]\n",
      " \n",
      "-from typing import Any, Mapping\n",
      "+import weakref\n",
      "+from typing import Any, Generic, Mapping, TypeVar, cast\n",
      "+\n",
      "+T = TypeVar('T')\n",
      " \n",
      " \n",
      " class StateDictMixin:\n",
      "@@ -19,4 +23,12 @@ class StateDictMixin:\n",
      "         pass\n",
      " \n",
      " \n",
      "-# TODO: define holder mixin\n",
      "+class HolderMixin(Generic[T]):\n",
      "+\n",
      "+    def __init__(self, *args, instance: T, **kwargs) -> None:\n",
      "+        super().__init__(*args, **kwargs)\n",
      "+        instance_proxy = (\n",
      "+            instance if isinstance(instance, weakref.ProxyTypes) else\n",
      "+            weakref.proxy(instance)\n",
      "+        )\n",
      "+        self._instance = cast(T, instance_proxy)\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"Validator\",\n",
    "    name=\"git_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    "    callbacks=[\n",
    "        dict(type=\"GitCallback\", diff='HEAD -- \":(exclude)*.ipynb\"'),\n",
    "    ],\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "\n",
    "    !echo\n",
    "    !cat {work_dirs}/git_callback/*.log"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:09,033 62058:140704275689088][base.py:56 todd.IterBasedTrainer.optimize_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:09,037 62058:140704275689088][log.py:93 todd.IterBasedTrainer.optimize_callback after_run_iter] INFO: Iter [1/8] batch={'x': tensor([4, 7]), 'y': tensor([ 8, 14])} weight=0.000 loss=11.000\n",
      "[2024-03-14 12:18:09,040 62058:140704275689088][log.py:93 todd.IterBasedTrainer.optimize_callback after_run_iter] INFO: Iter [2/8] batch={'x': tensor([6, 3]), 'y': tensor([12,  6])} weight=0.027 loss=8.876\n",
      "[2024-03-14 12:18:09,042 62058:140704275689088][log.py:93 todd.IterBasedTrainer.optimize_callback after_run_iter] INFO: Iter [3/8] batch={'x': tensor([1, 2]), 'y': tensor([2, 4])} weight=0.050 loss=2.925\n",
      "[2024-03-14 12:18:09,043 62058:140704275689088][log.py:93 todd.IterBasedTrainer.optimize_callback after_run_iter] INFO: Iter [4/8] batch={'x': tensor([ 5, 10]), 'y': tensor([10, 20])} weight=0.057 loss=14.569\n",
      "[2024-03-14 12:18:09,045 62058:140704275689088][log.py:93 todd.IterBasedTrainer.optimize_callback after_run_iter] INFO: Iter [5/8] batch={'x': tensor([8, 9]), 'y': tensor([16, 18])} weight=0.095 loss=16.193\n",
      "[2024-03-14 12:18:09,047 62058:140704275689088][log.py:93 todd.IterBasedTrainer.optimize_callback after_run_iter] INFO: Iter [6/8] batch={'x': tensor([7, 2]), 'y': tensor([14,  4])} weight=0.138 loss=8.381\n",
      "[2024-03-14 12:18:09,049 62058:140704275689088][log.py:93 todd.IterBasedTrainer.optimize_callback after_run_iter] INFO: Iter [7/8] batch={'x': tensor([ 8, 10]), 'y': tensor([16, 20])} weight=0.160 loss=16.560\n",
      "[2024-03-14 12:18:09,050 62058:140704275689088][log.py:93 todd.IterBasedTrainer.optimize_callback after_run_iter] INFO: Iter [8/8] batch={'x': tensor([9, 5]), 'y': tensor([18, 10])} weight=0.205 loss=12.565\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"IterBasedTrainer\",\n",
    "    name=\"optimize_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[\n",
    "        dict(type='OptimizeCallback'),\n",
    "        dict(type='LogCallback', interval=1),\n",
    "    ],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    iters=8,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Learning Rate Schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:09,063 62058:140704275689088][base.py:56 todd.IterBasedTrainer.lr_schedule_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:09,066 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [1/8] batch={'x': tensor([6, 2]), 'y': tensor([12,  4])} weight=0.000 loss=8.000 lr=['1.667e-03']\n",
      "[2024-03-14 12:18:09,068 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [2/8] batch={'x': tensor([ 9, 10]), 'y': tensor([18, 20])} weight=0.007 loss=18.937 lr=['2.333e-03']\n",
      "[2024-03-14 12:18:09,070 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [3/8] batch={'x': tensor([3, 8]), 'y': tensor([ 6, 16])} weight=0.029 loss=10.841 lr=['3.000e-03']\n",
      "[2024-03-14 12:18:09,073 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [4/8] batch={'x': tensor([5, 7]), 'y': tensor([10, 14])} weight=0.045 loss=11.728 lr=['3.667e-03']\n",
      "[2024-03-14 12:18:09,074 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [5/8] batch={'x': tensor([1, 4]), 'y': tensor([2, 8])} weight=0.067 loss=4.832 lr=['4.333e-03']\n",
      "[2024-03-14 12:18:09,077 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [6/8] batch={'x': tensor([1, 9]), 'y': tensor([ 2, 18])} weight=0.078 loss=9.609 lr=['5.000e-03']\n",
      "[2024-03-14 12:18:09,079 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [7/8] batch={'x': tensor([3, 2]), 'y': tensor([6, 4])} weight=0.103 loss=4.742 lr=['5.000e-03']\n",
      "[2024-03-14 12:18:09,080 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [8/8] batch={'x': tensor([7, 8]), 'y': tensor([14, 16])} weight=0.116 loss=14.132 lr=['5.000e-03']\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"IterBasedTrainer\",\n",
    "    name=\"lr_schedule_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[\n",
    "        dict(type=\"OptimizeCallback\"),\n",
    "        dict(\n",
    "            type=\"LRScheduleCallback\",\n",
    "            lr_scheduler=dict(type=\"LinearLR\", total_iters=5),\n",
    "        ),\n",
    "        dict(type=\"LogCallback\", interval=1),\n",
    "    ],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    iters=8,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:09,092 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.lr_schedule_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:09,094 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.lr_schedule_callback before_run_epoch] INFO: Epoch [1/5]\n",
      "[2024-03-14 12:18:09,096 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [1/10] batch={'x': tensor([3, 2]), 'y': tensor([6, 4])} weight=0.000 loss=5.000 lr=['1.667e-03']\n",
      "[2024-03-14 12:18:09,098 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [2/10] batch={'x': tensor([4, 1]), 'y': tensor([8, 2])} weight=0.004 loss=4.990 lr=['1.667e-03']\n",
      "[2024-03-14 12:18:09,099 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.lr_schedule_callback before_run_epoch] INFO: Epoch [2/5]\n",
      "[2024-03-14 12:18:09,101 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [3/10] batch={'x': tensor([4, 1]), 'y': tensor([8, 2])} weight=0.008 loss=4.979 lr=['2.778e-03']\n",
      "[2024-03-14 12:18:09,102 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [4/10] batch={'x': tensor([3, 2]), 'y': tensor([6, 4])} weight=0.015 loss=4.962 lr=['2.778e-03']\n",
      "[2024-03-14 12:18:09,104 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.lr_schedule_callback before_run_epoch] INFO: Epoch [3/5]\n",
      "[2024-03-14 12:18:09,106 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [5/10] batch={'x': tensor([1, 2]), 'y': tensor([2, 4])} weight=0.022 loss=2.967 lr=['3.889e-03']\n",
      "[2024-03-14 12:18:09,108 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [6/10] batch={'x': tensor([3, 4]), 'y': tensor([6, 8])} weight=0.028 loss=6.902 lr=['3.889e-03']\n",
      "[2024-03-14 12:18:09,109 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.lr_schedule_callback before_run_epoch] INFO: Epoch [4/5]\n",
      "[2024-03-14 12:18:09,112 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [7/10] batch={'x': tensor([1, 2]), 'y': tensor([2, 4])} weight=0.042 loss=2.938 lr=['5.000e-03']\n",
      "[2024-03-14 12:18:09,113 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [8/10] batch={'x': tensor([4, 3]), 'y': tensor([8, 6])} weight=0.049 loss=6.828 lr=['5.000e-03']\n",
      "[2024-03-14 12:18:09,114 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.lr_schedule_callback before_run_epoch] INFO: Epoch [5/5]\n",
      "[2024-03-14 12:18:09,117 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [9/10] batch={'x': tensor([2, 4]), 'y': tensor([4, 8])} weight=0.067 loss=5.800 lr=['5.000e-03']\n",
      "[2024-03-14 12:18:09,119 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.lr_schedule_callback after_run_iter] INFO: Iter [10/10] batch={'x': tensor([1, 3]), 'y': tensor([2, 6])} weight=0.082 loss=3.837 lr=['5.000e-03']\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"EpochBasedTrainer\",\n",
    "    name=\"lr_schedule_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=4),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[\n",
    "        dict(type=\"OptimizeCallback\"),\n",
    "        dict(\n",
    "            type=\"LRScheduleCallback\",\n",
    "            lr_scheduler=dict(type=\"LinearLR\", total_iters=3),\n",
    "            by_epoch=True,\n",
    "        ),\n",
    "        dict(type=\"LogCallback\", interval=1),\n",
    "    ],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    epochs=5,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Learning Rate Scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:18:09,131 62058:140704275689088][lr.py:93 todd.IterBasedTrainer.lr_scale_callback _scale_lr] INFO: base_batch_size=1 batch_size=2 lr_scaler=2.000\n",
      "\u001b[2m[2024-03-14 12:18:09,132 62058:140704275689088][base.py:56 todd.IterBasedTrainer.lr_scale_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:09,134 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_scale_callback after_run_iter] INFO: Iter [1/8] batch={'x': tensor([3, 4]), 'y': tensor([6, 8])} weight=0.000 loss=7.000\n",
      "[2024-03-14 12:18:09,136 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_scale_callback after_run_iter] INFO: Iter [2/8] batch={'x': tensor([6, 7]), 'y': tensor([12, 14])} weight=0.035 loss=12.773\n",
      "[2024-03-14 12:18:09,138 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_scale_callback after_run_iter] INFO: Iter [3/8] batch={'x': tensor([10,  2]), 'y': tensor([20,  4])} weight=0.100 loss=11.400\n",
      "[2024-03-14 12:18:09,140 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_scale_callback after_run_iter] INFO: Iter [4/8] batch={'x': tensor([8, 9]), 'y': tensor([16, 18])} weight=0.160 loss=15.640\n",
      "[2024-03-14 12:18:09,142 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_scale_callback after_run_iter] INFO: Iter [5/8] batch={'x': tensor([5, 1]), 'y': tensor([10,  2])} weight=0.245 loss=5.265\n",
      "[2024-03-14 12:18:09,145 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_scale_callback after_run_iter] INFO: Iter [6/8] batch={'x': tensor([6, 8]), 'y': tensor([12, 16])} weight=0.275 loss=12.075\n",
      "[2024-03-14 12:18:09,146 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_scale_callback after_run_iter] INFO: Iter [7/8] batch={'x': tensor([10,  3]), 'y': tensor([20,  6])} weight=0.345 loss=10.757\n",
      "[2024-03-14 12:18:09,148 62058:140704275689088][log.py:93 todd.IterBasedTrainer.lr_scale_callback after_run_iter] INFO: Iter [8/8] batch={'x': tensor([2, 4]), 'y': tensor([4, 8])} weight=0.410 loss=4.770\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"IterBasedTrainer\",\n",
    "    name=\"lr_scale_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[\n",
    "        dict(type=\"OptimizeCallback\"),\n",
    "        dict(\n",
    "            type=\"LRScaleCallback\",\n",
    "            lr_scaler=dict(base_batch_size=1),\n",
    "        ),\n",
    "        dict(type=\"LogCallback\", interval=1),\n",
    "    ],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    iters=8,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = todd.RunnerRegistry.build(\n",
    "        config,\n",
    "        work_dir=dict(root=work_dirs),\n",
    "    )\n",
    "    runner.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:09,174 62058:140704275689088][base.py:56 todd.IterBasedTrainer.checkpoint_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:09,177 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [1/8] batch={'x': tensor([6, 5]), 'y': tensor([12, 10])} weight=0.000 loss=11.000\n",
      "[2024-03-14 12:18:09,195 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_1\n",
      "[2024-03-14 12:18:09,224 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [2/8] batch={'x': tensor([2, 8]), 'y': tensor([ 4, 16])} weight=0.027 loss=9.863\n",
      "[2024-03-14 12:18:09,233 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_2\n",
      "[2024-03-14 12:18:09,241 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [3/8] batch={'x': tensor([ 3, 10]), 'y': tensor([ 6, 20])} weight=0.052 loss=12.659\n",
      "[2024-03-14 12:18:09,243 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_3\n",
      "[2024-03-14 12:18:09,251 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [4/8] batch={'x': tensor([1, 4]), 'y': tensor([2, 8])} weight=0.085 loss=4.787\n",
      "[2024-03-14 12:18:09,258 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_4\n",
      "[2024-03-14 12:18:09,262 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [5/8] batch={'x': tensor([9, 7]), 'y': tensor([18, 14])} weight=0.097 loss=15.220\n",
      "[2024-03-14 12:18:09,264 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_5\n",
      "[2024-03-14 12:18:09,268 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [6/8] batch={'x': tensor([9, 7]), 'y': tensor([18, 14])} weight=0.137 loss=14.900\n",
      "[2024-03-14 12:18:09,269 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_6\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:18:09,274 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [7/8] batch={'x': tensor([1, 3]), 'y': tensor([2, 6])} weight=0.177 loss=3.645\n",
      "[2024-03-14 12:18:09,276 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_7\n",
      "[2024-03-14 12:18:09,281 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [8/8] batch={'x': tensor([5, 8]), 'y': tensor([10, 16])} weight=0.187 loss=11.781\n",
      "[2024-03-14 12:18:09,282 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1;36m/var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460\u001b[0m\n",
      "└── \u001b[1;36mcheckpoint_callback\u001b[0m\n",
      "    └── \u001b[1;36mcheckpoints\u001b[0m\n",
      "        ├── \u001b[1;36miter_1\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_2\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_3\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_4\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_5\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_6\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_7\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_8\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        └── \u001b[35mlatest\u001b[0m -> \u001b[1;36miter_8\u001b[0m\n",
      "\n",
      "12 directories, 40 files\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:18:09,715 62058:140704275689088][checkpoint.py:54 todd.IterBasedTrainer.checkpoint_callback init] INFO: Loading from /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_5\n",
      "[2024-03-14 12:18:09,718 62058:140704275689088][base.py:65 todd.IterBasedTrainer.checkpoint_callback load_model_state_dict] INFO: <All keys matched successfully>\n",
      "\u001b[2m[2024-03-14 12:18:09,719 62058:140704275689088][base.py:56 todd.IterBasedTrainer.checkpoint_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:09,722 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [6/8] batch={'x': tensor([10,  1]), 'y': tensor([20,  2])} weight=0.137 loss=10.244\n",
      "[2024-03-14 12:18:09,723 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_6\n",
      "[2024-03-14 12:18:09,727 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [7/8] batch={'x': tensor([4, 2]), 'y': tensor([8, 4])} weight=0.165 loss=5.505\n",
      "[2024-03-14 12:18:09,728 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_7\n",
      "[2024-03-14 12:18:09,731 62058:140704275689088][log.py:93 todd.IterBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [8/8] batch={'x': tensor([6, 9]), 'y': tensor([12, 18])} weight=0.180 loss=13.650\n",
      "[2024-03-14 12:18:09,732 62058:140704275689088][checkpoint.py:80 todd.IterBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpkhq8r460/checkpoint_callback/checkpoints/iter_8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "strategy.pth:\n",
      "{}\n",
      "\n",
      "optim.pth:\n",
      "{'param_groups': [{'dampening': 0,\n",
      "                   'differentiable': False,\n",
      "                   'foreach': None,\n",
      "                   'lr': 0.005,\n",
      "                   'maximize': False,\n",
      "                   'momentum': 0,\n",
      "                   'nesterov': False,\n",
      "                   'params': [0],\n",
      "                   'weight_decay': 0}],\n",
      " 'state': {0: {'momentum_buffer': None}}}\n",
      "\n",
      "meta.pth:\n",
      "{'iter_': 5}\n",
      "\n",
      "model.pth:\n",
      "OrderedDict([('_weight', tensor(0.1375))])\n",
      "\n",
      "callbacks.pth:\n",
      "{'callbacks': [{}, {}, {}]}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"IterBasedTrainer\",\n",
    "    name=\"checkpoint_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[\n",
    "        dict(type='OptimizeCallback'),\n",
    "        dict(type='LogCallback', interval=1),\n",
    "        dict(type=\"CheckpointCallback\", interval=1),\n",
    "    ],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    iters=8,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.run()\n",
    "\n",
    "    !echo\n",
    "    !tree {work_dirs}\n",
    "    !echo\n",
    "\n",
    "    iter_5 = pathlib.Path(work_dirs) / 'checkpoint_callback' / 'checkpoints' / 'iter_5'\n",
    "    for f in iter_5.glob('*.pth'):\n",
    "        print(f'{f.name}:')\n",
    "        pprint(torch.load(f, 'cpu'))\n",
    "        print()\n",
    "\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(\n",
    "            config,\n",
    "            work_dir=dict(root=work_dirs),\n",
    "            load_from=str(iter_5),\n",
    "        )\n",
    "    runner.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:09,769 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.checkpoint_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:09,769 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [1/3]\n",
      "[2024-03-14 12:18:09,772 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [1/15] batch={'x': tensor([5, 3]), 'y': tensor([10,  6])} weight=0.000 loss=8.000\n",
      "[2024-03-14 12:18:09,774 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [2/15] batch={'x': tensor([2, 6]), 'y': tensor([ 4, 12])} weight=0.020 loss=7.920\n",
      "[2024-03-14 12:18:09,775 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_2\n",
      "[2024-03-14 12:18:09,779 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [3/15] batch={'x': tensor([4, 7]), 'y': tensor([ 8, 14])} weight=0.040 loss=10.780\n",
      "[2024-03-14 12:18:09,781 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [4/15] batch={'x': tensor([10,  1]), 'y': tensor([20,  2])} weight=0.067 loss=10.629\n",
      "[2024-03-14 12:18:09,782 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_4\n",
      "[2024-03-14 12:18:09,786 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [5/15] batch={'x': tensor([9, 8]), 'y': tensor([18, 16])} weight=0.095 loss=16.193\n",
      "[2024-03-14 12:18:09,787 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [2/3]\n",
      "[2024-03-14 12:18:09,789 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [6/15] batch={'x': tensor([5, 7]), 'y': tensor([10, 14])} weight=0.138 loss=11.175\n",
      "[2024-03-14 12:18:09,790 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_6\n",
      "[2024-03-14 12:18:09,794 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [7/15] batch={'x': tensor([1, 6]), 'y': tensor([ 2, 12])} weight=0.168 loss=6.414\n",
      "[2024-03-14 12:18:09,796 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [8/15] batch={'x': tensor([9, 3]), 'y': tensor([18,  6])} weight=0.185 loss=10.890\n",
      "[2024-03-14 12:18:09,797 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_8\n",
      "[2024-03-14 12:18:09,823 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [9/15] batch={'x': tensor([8, 2]), 'y': tensor([16,  4])} weight=0.215 loss=8.925\n",
      "[2024-03-14 12:18:09,825 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [10/15] batch={'x': tensor([10,  4]), 'y': tensor([20,  8])} weight=0.240 loss=12.320\n",
      "[2024-03-14 12:18:09,833 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_10\n",
      "[2024-03-14 12:18:09,842 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [3/3]\n",
      "[2024-03-14 12:18:09,844 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [11/15] batch={'x': tensor([5, 9]), 'y': tensor([10, 18])} weight=0.275 loss=12.075\n",
      "[2024-03-14 12:18:09,846 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [12/15] batch={'x': tensor([7, 2]), 'y': tensor([14,  4])} weight=0.310 loss=7.605\n",
      "[2024-03-14 12:18:09,857 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_12\n",
      "[2024-03-14 12:18:09,861 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [13/15] batch={'x': tensor([6, 3]), 'y': tensor([12,  6])} weight=0.333 loss=7.504\n",
      "[2024-03-14 12:18:09,863 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [14/15] batch={'x': tensor([ 1, 10]), 'y': tensor([ 2, 20])} weight=0.355 loss=9.048\n",
      "[2024-03-14 12:18:09,864 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_14\n",
      "[2024-03-14 12:18:09,868 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [15/15] batch={'x': tensor([8, 4]), 'y': tensor([16,  8])} weight=0.383 loss=9.705\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1;36m/var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9\u001b[0m\n",
      "└── \u001b[1;36mcheckpoint_callback\u001b[0m\n",
      "    └── \u001b[1;36mcheckpoints\u001b[0m\n",
      "        ├── \u001b[1;36miter_10\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_12\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_14\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_2\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_4\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_6\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36miter_8\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        └── \u001b[35mlatest\u001b[0m -> \u001b[1;36miter_14\u001b[0m\n",
      "\n",
      "11 directories, 35 files\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:18:10,289 62058:140704275689088][checkpoint.py:54 todd.EpochBasedTrainer.checkpoint_callback init] INFO: Loading from /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_8\n",
      "[2024-03-14 12:18:10,293 62058:140704275689088][base.py:65 todd.EpochBasedTrainer.checkpoint_callback load_model_state_dict] INFO: <All keys matched successfully>\n",
      "\u001b[2m[2024-03-14 12:18:10,294 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.checkpoint_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:10,295 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [2/3]\n",
      "[2024-03-14 12:18:10,298 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [9/15] batch={'x': tensor([3, 8]), 'y': tensor([ 6, 16])} weight=0.215 loss=9.818\n",
      "[2024-03-14 12:18:10,300 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [10/15] batch={'x': tensor([5, 4]), 'y': tensor([10,  8])} weight=0.243 loss=7.909\n",
      "[2024-03-14 12:18:10,300 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_10\n",
      "[2024-03-14 12:18:10,303 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [3/3]\n",
      "[2024-03-14 12:18:10,305 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [11/15] batch={'x': tensor([7, 4]), 'y': tensor([14,  8])} weight=0.265 loss=9.542\n",
      "[2024-03-14 12:18:10,307 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [12/15] batch={'x': tensor([10,  2]), 'y': tensor([20,  4])} weight=0.293 loss=10.245\n",
      "[2024-03-14 12:18:10,307 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_12\n",
      "[2024-03-14 12:18:10,311 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [13/15] batch={'x': tensor([6, 3]), 'y': tensor([12,  6])} weight=0.323 loss=7.549\n",
      "[2024-03-14 12:18:10,312 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [14/15] batch={'x': tensor([5, 9]), 'y': tensor([10, 18])} weight=0.345 loss=11.585\n",
      "[2024-03-14 12:18:10,313 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_14\n",
      "[2024-03-14 12:18:10,316 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [15/15] batch={'x': tensor([1, 8]), 'y': tensor([ 2, 16])} weight=0.380 loss=7.290\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--------------------\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:18:10,730 62058:140704275689088][checkpoint.py:54 todd.EpochBasedTrainer.checkpoint_callback init] INFO: Loading from /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_10\n",
      "[2024-03-14 12:18:10,734 62058:140704275689088][base.py:65 todd.EpochBasedTrainer.checkpoint_callback load_model_state_dict] INFO: <All keys matched successfully>\n",
      "\u001b[2m[2024-03-14 12:18:10,734 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.checkpoint_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:10,736 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [3/3]\n",
      "[2024-03-14 12:18:10,740 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [11/15] batch={'x': tensor([7, 6]), 'y': tensor([14, 12])} weight=0.265 loss=11.278\n",
      "[2024-03-14 12:18:10,742 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [12/15] batch={'x': tensor([1, 4]), 'y': tensor([2, 8])} weight=0.298 loss=4.256\n",
      "[2024-03-14 12:18:10,743 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_12\n",
      "[2024-03-14 12:18:10,748 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [13/15] batch={'x': tensor([8, 5]), 'y': tensor([16, 10])} weight=0.310 loss=10.985\n",
      "[2024-03-14 12:18:10,750 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [14/15] batch={'x': tensor([ 2, 10]), 'y': tensor([ 4, 20])} weight=0.343 loss=9.945\n",
      "[2024-03-14 12:18:10,751 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpn0pz4ax9/checkpoint_callback/checkpoints/iter_14\n",
      "[2024-03-14 12:18:10,756 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [15/15] batch={'x': tensor([9, 3]), 'y': tensor([18,  6])} weight=0.373 loss=9.765\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"EpochBasedTrainer\",\n",
    "    name=\"checkpoint_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[\n",
    "        dict(type='OptimizeCallback'),\n",
    "        dict(type='LogCallback', interval=1),\n",
    "        dict(type=\"CheckpointCallback\", interval=2),\n",
    "    ],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    epochs=3,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.run()\n",
    "\n",
    "    !echo\n",
    "    !tree {work_dirs}\n",
    "    !echo\n",
    "\n",
    "    iter_8 = pathlib.Path(work_dirs) / 'checkpoint_callback' / 'checkpoints' / 'iter_8'\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(\n",
    "            config,\n",
    "            work_dir=dict(root=work_dirs),\n",
    "            load_from=str(iter_8),\n",
    "        )\n",
    "    runner.run()\n",
    "\n",
    "    !echo\n",
    "    !echo {'-' * 20}\n",
    "    !echo\n",
    "\n",
    "    iter_10 = pathlib.Path(work_dirs) / 'checkpoint_callback' / 'checkpoints' / 'iter_10'\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(\n",
    "            config,\n",
    "            work_dir=dict(root=work_dirs),\n",
    "            load_from=str(iter_10),\n",
    "        )\n",
    "    runner.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:10,784 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.checkpoint_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:10,785 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [1/3]\n",
      "[2024-03-14 12:18:10,789 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [1/15] batch={'x': tensor([4, 5]), 'y': tensor([ 8, 10])} weight=0.000 loss=9.000\n",
      "[2024-03-14 12:18:10,793 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [2/15] batch={'x': tensor([3, 6]), 'y': tensor([ 6, 12])} weight=0.022 loss=8.899\n",
      "[2024-03-14 12:18:10,797 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [3/15] batch={'x': tensor([ 9, 10]), 'y': tensor([18, 20])} weight=0.045 loss=18.572\n",
      "[2024-03-14 12:18:10,799 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [4/15] batch={'x': tensor([8, 2]), 'y': tensor([16,  4])} weight=0.093 loss=9.538\n",
      "[2024-03-14 12:18:10,802 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [5/15] batch={'x': tensor([1, 7]), 'y': tensor([ 2, 14])} weight=0.117 loss=7.530\n",
      "[2024-03-14 12:18:10,803 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmplwjo14oz/checkpoint_callback/checkpoints/epoch_1\n",
      "[2024-03-14 12:18:10,808 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [2/3]\n",
      "[2024-03-14 12:18:10,812 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [6/15] batch={'x': tensor([4, 1]), 'y': tensor([8, 2])} weight=0.138 loss=4.656\n",
      "[2024-03-14 12:18:10,816 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [7/15] batch={'x': tensor([7, 5]), 'y': tensor([14, 10])} weight=0.150 loss=11.100\n",
      "[2024-03-14 12:18:10,820 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [8/15] batch={'x': tensor([3, 9]), 'y': tensor([ 6, 18])} weight=0.180 loss=10.920\n",
      "[2024-03-14 12:18:10,822 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [9/15] batch={'x': tensor([6, 2]), 'y': tensor([12,  4])} weight=0.210 loss=7.160\n",
      "[2024-03-14 12:18:10,824 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [10/15] batch={'x': tensor([10,  8]), 'y': tensor([20, 16])} weight=0.230 loss=15.930\n",
      "[2024-03-14 12:18:10,826 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmplwjo14oz/checkpoint_callback/checkpoints/epoch_2\n",
      "[2024-03-14 12:18:10,829 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [3/3]\n",
      "[2024-03-14 12:18:10,833 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [11/15] batch={'x': tensor([2, 9]), 'y': tensor([ 4, 18])} weight=0.275 loss=9.488\n",
      "[2024-03-14 12:18:10,835 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [12/15] batch={'x': tensor([6, 5]), 'y': tensor([12, 10])} weight=0.303 loss=9.336\n",
      "[2024-03-14 12:18:10,837 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [13/15] batch={'x': tensor([3, 8]), 'y': tensor([ 6, 16])} weight=0.330 loss=9.185\n",
      "[2024-03-14 12:18:10,839 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [14/15] batch={'x': tensor([ 4, 10]), 'y': tensor([ 8, 20])} weight=0.358 loss=11.497\n",
      "[2024-03-14 12:18:10,841 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [15/15] batch={'x': tensor([7, 1]), 'y': tensor([14,  2])} weight=0.393 loss=6.430\n",
      "[2024-03-14 12:18:10,843 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmplwjo14oz/checkpoint_callback/checkpoints/epoch_3\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1;36m/var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmplwjo14oz\u001b[0m\n",
      "└── \u001b[1;36mcheckpoint_callback\u001b[0m\n",
      "    └── \u001b[1;36mcheckpoints\u001b[0m\n",
      "        ├── \u001b[1;36mepoch_1\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36mepoch_2\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        ├── \u001b[1;36mepoch_3\u001b[0m\n",
      "        │   ├── callbacks.pth\n",
      "        │   ├── meta.pth\n",
      "        │   ├── model.pth\n",
      "        │   ├── optim.pth\n",
      "        │   └── strategy.pth\n",
      "        └── \u001b[35mlatest\u001b[0m -> \u001b[1;36mepoch_3\u001b[0m\n",
      "\n",
      "7 directories, 15 files\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-03-14 12:18:11,264 62058:140704275689088][checkpoint.py:54 todd.EpochBasedTrainer.checkpoint_callback init] INFO: Loading from /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmplwjo14oz/checkpoint_callback/checkpoints/epoch_2\n",
      "[2024-03-14 12:18:11,267 62058:140704275689088][base.py:65 todd.EpochBasedTrainer.checkpoint_callback load_model_state_dict] INFO: <All keys matched successfully>\n",
      "\u001b[2m[2024-03-14 12:18:11,268 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.checkpoint_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:11,269 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.checkpoint_callback before_run_epoch] INFO: Epoch [3/3]\n",
      "[2024-03-14 12:18:11,279 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [11/15] batch={'x': tensor([6, 2]), 'y': tensor([12,  4])} weight=0.275 loss=6.900\n",
      "[2024-03-14 12:18:11,291 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [12/15] batch={'x': tensor([3, 5]), 'y': tensor([ 6, 10])} weight=0.295 loss=6.820\n",
      "[2024-03-14 12:18:11,305 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [13/15] batch={'x': tensor([1, 4]), 'y': tensor([2, 8])} weight=0.315 loss=4.212\n",
      "[2024-03-14 12:18:11,312 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [14/15] batch={'x': tensor([7, 9]), 'y': tensor([14, 18])} weight=0.328 loss=13.380\n",
      "[2024-03-14 12:18:11,315 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.checkpoint_callback after_run_iter] INFO: Iter [15/15] batch={'x': tensor([10,  8]), 'y': tensor([20, 16])} weight=0.368 loss=14.693\n",
      "[2024-03-14 12:18:11,317 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.checkpoint_callback _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmplwjo14oz/checkpoint_callback/checkpoints/epoch_3\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"EpochBasedTrainer\",\n",
    "    name=\"checkpoint_callback\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[\n",
    "        dict(type='OptimizeCallback'),\n",
    "        dict(type='LogCallback', interval=1),\n",
    "        dict(type=\"CheckpointCallback\", interval=1, by_epoch=True),\n",
    "    ],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    epochs=3,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.run()\n",
    "\n",
    "    !echo\n",
    "    !tree {work_dirs}\n",
    "    !echo\n",
    "\n",
    "    epoch_2 = pathlib.Path(work_dirs) / 'checkpoint_callback' / 'checkpoints' / 'epoch_2'\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(\n",
    "            config,\n",
    "            work_dir=dict(root=work_dirs),\n",
    "            load_from=str(epoch_2),\n",
    "        )\n",
    "    runner.run()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Monitor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomError(RuntimeError):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "@todd.RunnerRegistry.register_()\n",
    "class FaultyValidator(todd.runners.Validator):\n",
    "\n",
    "    def _run_iter(self, *args, **kwargs) -> NoReturn:\n",
    "        raise CustomError(\"faulty runner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:11,358 62058:140704275689088][base.py:56 todd.FaultyValidator.monitor_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "\u001b[1;31m[2024-03-14 12:18:11,359 62058:140704275689088][monitor.py:26 todd.FaultyValidator.monitor_callback __exit__] ERROR: Unable to run iter_=1\n",
      "batch={'x': tensor([1]), 'y': tensor([2])}\n",
      "memo={'dataloader': <torch.utils.data.dataloader.DataLoader object at 0x152023210>}\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/bytedance/.local/share/virtualenvs/todd-ARrcnwyq/lib/python3.11/site-packages/todd/runners/base.py\", line 246, in _run\n",
      "    memo = self._run_iter(batch, memo)\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/ipykernel_62058/1715875531.py\", line 5, in _run_iter\n",
      "    raise CustomError(\"faulty runner\")\n",
      "CustomError: faulty runner\u001b[m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[2024-03-14 12:18:11,358 62058:140704275689088][base.py:56 todd.FaultyValidator.monitor_callback __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\n",
      "[2024-03-14 12:18:11,359 62058:140704275689088][monitor.py:26 todd.FaultyValidator.monitor_callback __exit__] ERROR: Unable to run iter_=1\n",
      "batch={'x': tensor([1]), 'y': tensor([2])}\n",
      "memo={'dataloader': <torch.utils.data.dataloader.DataLoader object at 0x152023210>}\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/bytedance/.local/share/virtualenvs/todd-ARrcnwyq/lib/python3.11/site-packages/todd/runners/base.py\", line 246, in _run\n",
      "    memo = self._run_iter(batch, memo)\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/ipykernel_62058/1715875531.py\", line 5, in _run_iter\n",
      "    raise CustomError(\"faulty runner\")\n",
      "CustomError: faulty runner\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type='FaultyValidator',\n",
    "    name='monitor_callback',\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=20),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=1),\n",
    "    callbacks=[\n",
    "        dict(type='MonitorCallback'),\n",
    "        dict(type='LogCallback', interval=5, with_file_handler=True),\n",
    "    ],\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    try:\n",
    "        runner.run()\n",
    "    except CustomError as e:\n",
    "        pass\n",
    "\n",
    "    !echo\n",
    "    !cat {work_dirs}/monitor_callback/*.log\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Priorities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Strategies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:11,674 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.strategy_load_model_from __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:11,675 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.strategy_load_model_from before_run_epoch] INFO: Epoch [1/3]\n",
      "[2024-03-14 12:18:11,679 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [1/15] batch={'x': tensor([2, 5]), 'y': tensor([ 4, 10])} weight=0.000 loss=7.000\n",
      "[2024-03-14 12:18:11,682 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [2/15] batch={'x': tensor([ 7, 10]), 'y': tensor([14, 20])} weight=0.018 loss=16.851\n",
      "[2024-03-14 12:18:11,684 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [3/15] batch={'x': tensor([1, 3]), 'y': tensor([2, 6])} weight=0.060 loss=3.880\n",
      "[2024-03-14 12:18:11,687 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [4/15] batch={'x': tensor([8, 4]), 'y': tensor([16,  8])} weight=0.070 loss=11.580\n",
      "[2024-03-14 12:18:11,689 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [5/15] batch={'x': tensor([6, 9]), 'y': tensor([12, 18])} weight=0.100 loss=14.250\n",
      "[2024-03-14 12:18:11,691 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.strategy_load_model_from _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpj2udcidm/strategy_load_model_from/checkpoints/epoch_1\n",
      "[2024-03-14 12:18:11,694 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.strategy_load_model_from before_run_epoch] INFO: Epoch [2/3]\n",
      "[2024-03-14 12:18:11,697 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [6/15] batch={'x': tensor([9, 4]), 'y': tensor([18,  8])} weight=0.138 loss=12.106\n",
      "[2024-03-14 12:18:11,700 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [7/15] batch={'x': tensor([ 5, 10]), 'y': tensor([10, 20])} weight=0.170 loss=13.725\n",
      "[2024-03-14 12:18:11,702 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [8/15] batch={'x': tensor([3, 2]), 'y': tensor([6, 4])} weight=0.207 loss=4.481\n",
      "[2024-03-14 12:18:11,704 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [9/15] batch={'x': tensor([8, 6]), 'y': tensor([16, 12])} weight=0.220 loss=12.460\n",
      "[2024-03-14 12:18:11,706 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [10/15] batch={'x': tensor([7, 1]), 'y': tensor([14,  2])} weight=0.255 loss=6.980\n",
      "[2024-03-14 12:18:11,708 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.strategy_load_model_from _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpj2udcidm/strategy_load_model_from/checkpoints/epoch_2\n",
      "[2024-03-14 12:18:11,711 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.strategy_load_model_from before_run_epoch] INFO: Epoch [3/3]\n",
      "[2024-03-14 12:18:11,713 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [11/15] batch={'x': tensor([6, 1]), 'y': tensor([12,  2])} weight=0.275 loss=6.038\n",
      "[2024-03-14 12:18:11,714 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [12/15] batch={'x': tensor([7, 9]), 'y': tensor([14, 18])} weight=0.293 loss=13.660\n",
      "[2024-03-14 12:18:11,716 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [13/15] batch={'x': tensor([8, 5]), 'y': tensor([16, 10])} weight=0.333 loss=10.839\n",
      "[2024-03-14 12:18:11,718 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [14/15] batch={'x': tensor([ 4, 10]), 'y': tensor([ 8, 20])} weight=0.365 loss=11.445\n",
      "[2024-03-14 12:18:11,720 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [15/15] batch={'x': tensor([2, 3]), 'y': tensor([4, 6])} weight=0.400 loss=4.000\n",
      "[2024-03-14 12:18:11,721 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.strategy_load_model_from _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpj2udcidm/strategy_load_model_from/checkpoints/epoch_3\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--------------------\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m[2024-03-14 12:18:12,135 62058:140704275689088][base.py:56 todd.EpochBasedTrainer.strategy_load_model_from __init__] DEBUG: Rank 0 initialized by bytedance@C02G870SMD6R\u001b[m\n",
      "[2024-03-14 12:18:12,135 62058:140704275689088][base.py:80 todd.EpochBasedTrainer.strategy_load_model_from load_model_from] INFO: Loading model from /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpj2udcidm/strategy_load_model_from/checkpoints/epoch_2/model.pth\n",
      "[2024-03-14 12:18:12,138 62058:140704275689088][base.py:65 todd.EpochBasedTrainer.strategy_load_model_from load_model_state_dict] INFO: <All keys matched successfully>\n",
      "[2024-03-14 12:18:12,139 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.strategy_load_model_from before_run_epoch] INFO: Epoch [1/3]\n",
      "[2024-03-14 12:18:12,142 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [1/15] batch={'x': tensor([10,  1]), 'y': tensor([20,  2])} weight=0.275 loss=9.488\n",
      "[2024-03-14 12:18:12,143 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [2/15] batch={'x': tensor([5, 3]), 'y': tensor([10,  6])} weight=0.303 loss=6.790\n",
      "[2024-03-14 12:18:12,145 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [3/15] batch={'x': tensor([9, 6]), 'y': tensor([18, 12])} weight=0.323 loss=12.581\n",
      "[2024-03-14 12:18:12,146 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [4/15] batch={'x': tensor([2, 4]), 'y': tensor([4, 8])} weight=0.360 loss=4.920\n",
      "[2024-03-14 12:18:12,148 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [5/15] batch={'x': tensor([8, 7]), 'y': tensor([16, 14])} weight=0.375 loss=12.188\n",
      "[2024-03-14 12:18:12,149 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.strategy_load_model_from _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpj2udcidm/strategy_load_model_from/checkpoints/epoch_1\n",
      "[2024-03-14 12:18:12,153 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.strategy_load_model_from before_run_epoch] INFO: Epoch [2/3]\n",
      "[2024-03-14 12:18:12,154 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [6/15] batch={'x': tensor([8, 1]), 'y': tensor([16,  2])} weight=0.412 loss=7.144\n",
      "[2024-03-14 12:18:12,156 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [7/15] batch={'x': tensor([9, 3]), 'y': tensor([18,  6])} weight=0.435 loss=9.390\n",
      "[2024-03-14 12:18:12,157 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [8/15] batch={'x': tensor([10,  4]), 'y': tensor([20,  8])} weight=0.465 loss=10.745\n",
      "[2024-03-14 12:18:12,159 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [9/15] batch={'x': tensor([2, 6]), 'y': tensor([ 4, 12])} weight=0.500 loss=6.000\n",
      "[2024-03-14 12:18:12,161 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [10/15] batch={'x': tensor([5, 7]), 'y': tensor([10, 14])} weight=0.520 loss=8.880\n",
      "[2024-03-14 12:18:12,162 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.strategy_load_model_from _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpj2udcidm/strategy_load_model_from/checkpoints/epoch_2\n",
      "[2024-03-14 12:18:12,164 62058:140704275689088][log.py:99 todd.EpochBasedTrainer.strategy_load_model_from before_run_epoch] INFO: Epoch [3/3]\n",
      "[2024-03-14 12:18:12,166 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [11/15] batch={'x': tensor([ 6, 10]), 'y': tensor([12, 20])} weight=0.550 loss=11.600\n",
      "[2024-03-14 12:18:12,167 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [12/15] batch={'x': tensor([7, 1]), 'y': tensor([14,  2])} weight=0.590 loss=5.640\n",
      "[2024-03-14 12:18:12,169 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [13/15] batch={'x': tensor([9, 3]), 'y': tensor([18,  6])} weight=0.610 loss=8.340\n",
      "[2024-03-14 12:18:12,170 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [14/15] batch={'x': tensor([4, 2]), 'y': tensor([8, 4])} weight=0.640 loss=4.080\n",
      "[2024-03-14 12:18:12,172 62058:140704275689088][log.py:93 todd.EpochBasedTrainer.strategy_load_model_from after_run_iter] INFO: Iter [15/15] batch={'x': tensor([5, 8]), 'y': tensor([10, 16])} weight=0.655 loss=8.743\n",
      "[2024-03-14 12:18:12,173 62058:140704275689088][checkpoint.py:80 todd.EpochBasedTrainer.strategy_load_model_from _save] INFO: Saving state dict to /var/folders/v_/1kkfntxs5z74_rwvy1f3_mp80000gn/T/tmpj2udcidm/strategy_load_model_from/checkpoints/epoch_3\n"
     ]
    }
   ],
   "source": [
    "config = todd.Config(\n",
    "    type=\"EpochBasedTrainer\",\n",
    "    name=\"strategy_load_model_from\",\n",
    "    strategy=dict(type='BaseStrategy'),\n",
    "    dataset=dict(type='RunnerDataset', n=10),\n",
    "    model=dict(type='RunnerModel'),\n",
    "    dataloader=dict(batch_size=2, shuffle=True),\n",
    "    callbacks=[\n",
    "        dict(type='OptimizeCallback'),\n",
    "        dict(type='LogCallback', interval=1),\n",
    "        dict(type=\"CheckpointCallback\", interval=1, by_epoch=True),\n",
    "    ],\n",
    "    optimizer=dict(type=\"SGD\", lr=0.005),\n",
    "    epochs=3,\n",
    ")\n",
    "with tempfile.TemporaryDirectory() as work_dirs:\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.run()\n",
    "\n",
    "    !echo\n",
    "    !echo {'-' * 20}\n",
    "    !echo\n",
    "\n",
    "    epoch_2 = (pathlib.Path(work_dirs) / 'strategy_load_model_from' / 'checkpoints' / 'epoch_2' / 'model.pth')\n",
    "    runner: todd.runners.BaseRunner = \\\n",
    "        todd.RunnerRegistry.build(config, work_dir=dict(root=work_dirs))\n",
    "    runner.strategy.load_model_from(epoch_2)\n",
    "    runner.run()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dry Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "todd.Store.DRY_RUN = True"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "todd",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "0fe19504897982c0d86de0bd38ea30a541b47032e25039ac5ae6cd1de5b1a414"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
